#include <malloc.h>
#include <process.h>
#include <stdarg.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <winsock2.h>

#include "lib.h"

#include "growbuf.h"

#include "log.h"
#include "smtp.h"

#define SMTP_PORT	25


class SMTPWorker
{
	public:
		SMTPWorker(const String& mailhost, const SMTP::Header&, const String& message);
		~SMTPWorker();

		bool8 start();
		bool8 run();

	protected:


		enum ServerStates
		{
			STATE_WAIT_GREETING,
			STATE_HELO,
			STATE_MAIL_FROM,
			STATE_RCPT_TO,
			STATE_DATA,
			STATE_MESSAGE,
		};

		static DWORD __stdcall threadProc(SMTPWorker*);

		bool8 connect();
		bool8 prelude();
		bool8 sendHeader();
		bool8 sendMessage();
		bool8 sendQuit();
		void disconnect();

	private:
		bool8 send(int32 nArgs, ...);
		bool8 getResponse();
		bool8 sendRCPTs(const char* command, const char* headerList, bool8 bRequired);
		String formatHeaderLine(char* unformatted);

	protected:
		enum Constants { WORKER_BUFSIZE = 256 };

		SMTP::ASCIIHeader header;
		String mailhost;
		String msg;
		SOCKET sock;
		ServerStates state;
		char buf[WORKER_BUFSIZE];
		static const char*	m_delim;
};

const char* SMTPWorker::m_delim = " \t\n\r;,";	// delimeter list

//send message with utf-8 encoding
bool8
SMTP::send(const String& mailhost, const Header& hdr, const String& message)
{
	if (!mailhost.isEmpty() && !hdr.emailFrom.isEmpty() 
			&& !hdr.emailTo.isEmpty())
	{
		SMTPWorker* worker = new SMTPWorker(mailhost, hdr, message);

		if (worker && worker->start())
			return true;	// worker will kill itself upon success or failure
		else
			delete worker;
	}

	return false;
}

//send message with default encoding
bool8
SMTP::sendDefault(const String& mailhost, const Header& inhdr, const String& message)
{
	Header hdr(inhdr);
	hdr.defaultCodepage = true;
	if (!mailhost.isEmpty() && !hdr.emailFrom.isEmpty() 
			&& !hdr.emailTo.isEmpty())
	{
		SMTPWorker* worker = new SMTPWorker(mailhost, hdr, message);

		if (worker && worker->start())
			return true;	// worker will kill itself upon success or failure
		else
			delete worker;
	}

	return false;
}


bool8
SMTP::blockingSend(const String& mailhost, const Header& hdr, const String& message)
{
	bool8 result = false;

	if (!mailhost.isEmpty() && !hdr.emailFrom.isEmpty() 
			&& !hdr.emailTo.isEmpty())
	{
		SMTPWorker* worker = new SMTPWorker(mailhost, hdr, message);

		if (worker && worker->run())
			result =  true;

		delete worker;
	}

	return result;
}


//blocking send with default encoding
bool8
SMTP::blockingSendDefault(const String& mailhost, const Header& inhdr, const String& message)
{
	bool8 result = false;

	Header hdr(inhdr);
	hdr.defaultCodepage = true;
	if (!mailhost.isEmpty() && !hdr.emailFrom.isEmpty() 
			&& !hdr.emailTo.isEmpty())
	{
		SMTPWorker* worker = new SMTPWorker(mailhost, hdr, message);

		if (worker && worker->run())
			result =  true;

		delete worker;
	}

	return result;
}



SMTPWorker::SMTPWorker(const String& mailhost, const SMTP::Header& hdr, const String& message)
{
	header = hdr;
	this->mailhost = mailhost;
	msg = message;
	sock = INVALID_SOCKET;
}

SMTPWorker::~SMTPWorker()
{
	disconnect();
}

bool8
SMTPWorker::start()
{
	HANDLE hThread;
	DWORD threadID;

	if (!mailhost.isEmpty() && 
		header.emailFrom[0] && 
		header.emailTo[0])
	{
		hThread = (HANDLE)_beginthreadex(
			(void*)0, 0, (unsigned(__stdcall*)(void*))threadProc,
			(void*)this, 0, (unsigned*)&threadID);
		if (hThread)
		{
			CloseHandle(hThread);
			return true;
		}
	}
	return false;	// caller must delete this SMTPWorker!
}

DWORD __stdcall
SMTPWorker::threadProc(SMTPWorker* This)
{
	DWORD result = 0;

	result = This->run();

	delete This;

	_endthreadex(result);
	return result;
}

bool8
SMTPWorker::run()
{
	bool8 result = false;

	ASSERT(
		!mailhost.isEmpty() &&
		header.emailFrom[0] &&
		header.emailTo[0]);

	if (connect())
	{		
		result =
			prelude() &&
			sendHeader() &&
			sendMessage() &&
			sendQuit();
		
		disconnect();
	}
	else
	{
		Log::msgError(TEXT("Failed to connect to mail host: '%s'"), 
			(const uint16*)mailhost);
	}

	return result;
}





bool8
SMTPWorker::connect()
{
	SOCKADDR_IN addr;
	uint32 ipaddr;

	WORD version = MAKEWORD(2,0);
	WSADATA wsaData;
	int err = WSAStartup( version, &wsaData );
	if ( err != 0 ) 
	{
		return false;
	}

	if ( LOBYTE( wsaData.wVersion ) != 2 ||
         HIBYTE( wsaData.wVersion ) != 0 ) 
	{
		WSACleanup();
		return false;
	}
	addr.sin_family = AF_INET;
	addr.sin_port = htons(SMTP_PORT);

	int32 len = mailhost.getLength();
	char* hostname = (char*)_alloca(len + 1);
	StringConversions::convert(hostname, mailhost, len + 1);
	
	ipaddr = ntohl(inet_addr(hostname));

	if (ipaddr == INADDR_NONE)
	{
		HOSTENT* host = gethostbyname(hostname);
		if (host)
			ipaddr = ntohl(((IN_ADDR*)host->h_addr_list[0])->s_addr);
		else
		{
			return false;
		}
	}

	addr.sin_addr.s_addr = htonl(ipaddr);

	sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (sock == INVALID_SOCKET)
		return false;

	if (SOCKET_ERROR == ::connect(sock, (PSOCKADDR)&addr, sizeof(SOCKADDR_IN)))
	{
		closesocket(sock);
		sock = INVALID_SOCKET;
		return false;
	}

	return true;
}

// this function uses char buffers as parameters
// do not use Strings
bool8
SMTPWorker::send(int32 nArgs, ...)
{
	va_list args;
	WSABUF* scatter;
	int32 i, j;
	bool8 result = false;
	DWORD total, wasSent;

	va_start(args, nArgs);

	scatter = new WSABUF[nArgs];
	if (scatter)
	{
		total = 0;

		for (i = j = 0; i < nArgs; i++)
		{
			scatter[j].buf = va_arg(args, char*);
			if (scatter[j].buf)
			{
				scatter[j].len = strlen(scatter[j].buf);
				if (scatter[j].len > 0)
				{
					total += scatter[j].len;
					++j;
				}
			}
		}

		if (j > 0)
		{
			result = (bool8)(
				0 == WSASend(sock, scatter, j,  &wasSent, 0,
					(LPWSAOVERLAPPED)0,
					(LPWSAOVERLAPPED_COMPLETION_ROUTINE)0) &&
				wasSent == total);
		}
		else
			result = true;

		delete [] scatter;
	}

	va_end(args);

	return result;
}

bool8
SMTPWorker::getResponse()
{
	int result = recv(sock, buf, WORKER_BUFSIZE, 0);
	
	if (result == SOCKET_ERROR)
		return false;

	buf[result] = '\0';

	switch (state)
	{
		case STATE_WAIT_GREETING:
			result = (strncmp(buf, "220", 3) == 0);
			break;
		
		case STATE_HELO:
		case STATE_MAIL_FROM:
		case STATE_MESSAGE:
			result = (strncmp(buf, "250", 3) == 0);
			break;

		case STATE_RCPT_TO:
			result = (strncmp(buf, "250", 3) == 0 ||
					  strncmp(buf, "251", 3) == 0);
			break;

		case STATE_DATA:
			result = (strncmp(buf, "354", 3) == 0);
			break;

		default:
			result = false;
	}

	if (!result)
	{
		char* strstate;

		switch (state)
		{
			case STATE_WAIT_GREETING:
				strstate = "STATE_WAIT_GREETING";
				break;
			
			case STATE_HELO:
				strstate = "STATE_HELO";
				break;

			case STATE_MAIL_FROM:
				strstate = "STATE_MAIL_FROM";
				break;

			case STATE_RCPT_TO:
				strstate = "STATE_RCPT_TO";
				break;

			case STATE_MESSAGE:
				strstate = "STATE_MESSAGE";
				break;

			case STATE_DATA:
				strstate = "STATE_DATA";
				break;

			default:
				strstate = "STATE_UNKNOWN";
		}


		Log::msgError("SMTPWorker::getResponse() invalid response: '%s' for state '%s'.", 
			buf, strstate);
	}

	return (bool8)result;

}

bool8 
SMTPWorker::prelude()
{
	char localhost[256];
	bool8 result;

	ASSERT(sock != INVALID_SOCKET);

	state = STATE_WAIT_GREETING;

	result = getResponse();

	if (result)
	{
		state = STATE_HELO;

		gethostname(localhost, 256);
		result = send(3, "HELO ", localhost, "\r\n") && getResponse();
	}

	if (result)
	{
		state = STATE_MAIL_FROM;
		result = send(3, "MAIL FROM: <", header.emailFrom, ">\r\n") && getResponse();
	}
	
	if (result)
	{
		state = STATE_RCPT_TO;
		result = sendRCPTs("RCPT TO:", header.emailTo, true);
	}


	if (result)
	{
		result = sendRCPTs("RCPT TO:", header.emailCC, false);
	}


	if (result)
	{
		result = sendRCPTs("RCPT TO:", header.emailBCC, false);
	}

	return result;
}



/*
 *  This function adds support for multiple email addresses
 *  seperated by the commas, semicolons and whitespace
 */
bool8
SMTPWorker::sendRCPTs(const char* command, const char* headerList, bool8 bRequired)
{
	bool8				result		  = true;
	int					emailListSize = 0;	        // in bytes
	char*				emailList	  = NULL;       // stores header.emailTo
	char*				email		  = NULL;       // found email address
	char				cmd[32];			        // ASCII version of command
	unsigned int		count		  = 0;			// ensure data was sent

	// Format the command in ASCII, as SMTP is an ASCII protocol
	_snprintf(cmd, sizeof(cmd), "%s<", command);

	// The email addresses are stored in headerList.
	// Copy the email addresses to a local buffer so we do not destroy
	// the global header.
	emailListSize = strlen(headerList) + 1;
	emailList = (char *) _alloca(emailListSize);
	memcpy(emailList, headerList, emailListSize);
	
	email = strtok(emailList, m_delim);
	while (email)
	{
		result = send(3, cmd, email, ">\r\n") && getResponse();
		
		if (!result)
			break;
		
		count++;
		emailList = NULL;
		email = strtok(emailList, m_delim);
	}

	if (!count && bRequired)
		result = false;

	return result;
}


/*
 * Creates a formatted list of email addresses as specified in RFC 822. 
 *
 * Parameters:
 * 'unformatted' - an unformatted list of email addresses that is
 *                 at least separated by the characters composed of 'delim'.
 *
 * Return value:  Returns the formatted list of email addresses separated by commas.
 */
String
SMTPWorker::formatHeaderLine(char* unformatted)
{
	char*				email		= NULL;			// stripped email
	String				formatted	= "";

	email = strtok(unformatted, m_delim);
	while (email)
	{
		formatted += email;
		formatted += ", ";
		unformatted = NULL;
		email = strtok(unformatted, m_delim);	
	}

	formatted.setChar(formatted.getLength() - 2, TEXT('\0'));
	
	return formatted;
}



bool8
SMTPWorker::sendHeader()
{
	bool8	result;
	String	strHeaderLine;
	char*	headerLine		= NULL;

	ASSERT(sock != INVALID_SOCKET);

	state = STATE_DATA;

	result = send(1, "DATA\r\n") && getResponse();

	if (result)
	{
		strHeaderLine = formatHeaderLine(header.emailTo);
		headerLine = (char*) _alloca(strHeaderLine.getLength() + 1);
		StringConversions::convert(headerLine, strHeaderLine, strHeaderLine.getLength() + 1);

		result = send(3, "To: ", headerLine, "\r\n");
	}

	if (result)
	{
		strHeaderLine = formatHeaderLine(header.emailCC);
		headerLine = (char*) _alloca(strHeaderLine.getLength() + 1);
		StringConversions::convert(headerLine, strHeaderLine, strHeaderLine.getLength() + 1);

		result = send(3, "CC: ", headerLine, "\r\n");
	}


	if (result)
	{
		if (header.nameFrom[0])
			result = send(5, "From: ", header.nameFrom,
				" <", header.emailFrom, ">\r\n");
		else
			result = send(3, "From: ", header.emailFrom, "\r\n");
	}
	
	if (result && header.subject[0])
		result = send(3, "Subject: ", header.subject, "\r\n");

	if (result)
		result = send(1, "MIME-Version: 1.0\r\n");
	
	if (result && header.bEmailAsHtml)
	{
		if (!header.defaultCodepage)
			result = send(1, "Content-Type: text/html; charset=utf-8\r\n");
		else
			result = send(1, "Content-Type: text/html; charset=\"DEFAULT\"\r\n");
	}
	else if (result && !header.bEmailAsHtml)
	{
		if (!header.defaultCodepage)
			result = send(1, "Content-Type: text/plain; charset=utf-8\r\n");
		else
			result = send(1, "Content-Type: text/plain; charset=\"DEFAULT\"\r\n");
	}
	
	if (result)
		result = send(1, "Content-Transfer-Encoding: Quoted-Printable\r\n");

	if (result)
		result = send(1, "\r\n");	// separator between header and message body

	return result;
}

// Simple function to convert the lower 4 bits of a byte into its
// ASCII hex character representation
inline uint8 b2ah(uint8 i)
{
	if (i >= 0 && i <= 9)
		return i + 48;
	else if (i >= 10 && i <= 15)
		return i + 55;

	return 0;
}

bool8
SMTPWorker::sendMessage()
{
	ASSERT(sock != INVALID_SOCKET);

	state = STATE_MESSAGE;

	//max size of utf-8 is 6 characters per unicode char
	//max size of html entity is 8 characters per unicode car
	uint32 utf8len = msg.getLength() * 8 + 1;
	char* utf8msg = new char [utf8len];

	if (!utf8msg)
		return false;

	*utf8msg = '\0';

	if (header.bEmailAsHtml)
	{
		/* converting unicode values to html decimal enities so they show up no matter 
			what the encoding is set to */
		for(uint32 i=0,j=0;i<msg.getLength();i++,j++)
		{
			if (msg[i] < 0x7f)
				utf8msg[j] = (char)msg[i];
			else
			{
				int wrote = sprintf(&utf8msg[j], "&#%d;", msg[i]);
				if (wrote > 0)
					j+= (wrote - 1);
			}
		}
		utf8msg[j] = '\0';
	}
	else
	{
		StringConversions::encodeUTF8(utf8msg, msg, utf8len);
	}
	
	GrowBuf safemsg;
	uint8 line[76 + 3];
	uint8* inmsg = (uint8*)utf8msg;
	int32 linecount = 0;
	while(inmsg && *inmsg)
	{
		/* 
		In particular, the only characters that are known to be 
		consistent across all gateways are the 73 characters that 
		correspond to the upper and lower case letters A-Z and a-z, 
		the 10 digits 0-9, and the following eleven special characters: 
						"'"  (ASCII code 39)
						"("  (ASCII code 40)
						")"  (ASCII code 41)
						"+"  (ASCII code 43)
						","  (ASCII code 44)
						"-"  (ASCII code 45)
						"."  (ASCII code 46)
						"/"  (ASCII code 47)
						":"  (ASCII code 58)
						"="  (ASCII code 61)
						"?"  (ASCII code 63)
		*/
 		if ((*inmsg >= 'A' && *inmsg <= 'Z') || (*inmsg >= 'a' && *inmsg <= 'z') ||
			(*inmsg >= '0' && *inmsg <= '9') || *inmsg == '\'' || *inmsg == '(' ||
			*inmsg == ')' || *inmsg == '+' ||	*inmsg == ',' || *inmsg == '-' ||	
			*inmsg == '.' || *inmsg== '/' || *inmsg == ':' || *inmsg == '?'
			)
		{
			// 76 max line count
			if (linecount >= 76 - 4)
			{
				line[linecount++] = '=';
				line[linecount] = '\r';
				line[linecount + 1] = '\n';

				line[linecount + 2] ='\0';
				safemsg.add((char*)line);

				ASSERT( ++linecount <= 76);
				
				linecount = 0;
			}

			//it's safe
			line[linecount++]  = *inmsg;
		}
		else
			// octet must be used for characters > 127, =, SPACE, or TAB,
			// and is recommended for any characters not listed above
		{
			// 76 max line count
			if (linecount >= 76 - 4)
			{
				line[linecount++] = '=';
				line[linecount] = '\r';
				line[linecount + 1] = '\n';

				line[linecount + 2] ='\0';
				safemsg.add((char*)line);

				ASSERT( ++linecount <= 76);

				linecount = 0;
			}

			line[linecount++]  = '=';
			line[linecount++]  = b2ah((*inmsg >> 4) & 0xF);
			line[linecount++]  = b2ah(*inmsg & 0xF);
		}

		inmsg++;
	}
	
	if (linecount)
	{
		line[linecount] ='\0';
		safemsg.add((char*)line);
	}
	
	bool8 result = send(2, safemsg.getBuf(), "\r\n.\r\n") && getResponse();
	
	delete[] utf8msg;
	utf8msg = 0;

	return result;
}

bool8
SMTPWorker::sendQuit()
{
	ASSERT(sock != INVALID_SOCKET);

	return send(1, "QUIT\r\n");
}

void
SMTPWorker::disconnect()
{
	if (sock != INVALID_SOCKET)
	{
		closesocket(sock);
		sock = INVALID_SOCKET;
	}
}


